# bench/benchmark.py
import asyncio
import time
import random
import statistics
import matplotlib.pyplot as plt
import aiohttp

SERVER = "http://localhost:8000"
ENDPOINT = "/v1/generate"
MODEL = "facebook/opt-6.7b"  # adjust to available model / local path
PROMPTS = [
    "Write a short poem about the ocean.",
    "Summarize the following paragraph: ...",
    "Explain reinforcement learning in simple terms.",
]

async def single_client(session, rps, duration, results):
    interval = 1.0 / rps
    endt = time.time() + duration
    while time.time() < endt:
        t0 = time.time()
        data = {"model": MODEL, "prompt": random.choice(PROMPTS)}
        try:
            async with session.post(SERVER + ENDPOINT, json=data, timeout=10) as resp:
                t1 = time.time()
                if resp.status == 200:
                    results.append((True, t1 - t0))
                else:
                    results.append((False, t1 - t0))
        except Exception as e:
            results.append((False, time.time() - t0))
        # sleep to maintain rate
        now = time.time()
        to_wait = interval - (now - t0)
        if to_wait > 0:
            await asyncio.sleep(to_wait)

async def run_rps_point(rps, clients, duration=20):
    # spawn `clients` coroutines each generating rps/clients requests per second (so total RPS ~ rps)
    results = []
    async with aiohttp.ClientSession() as session:
        tasks = []
        per_client = rps / clients
        for _ in range(clients):
            tasks.append(asyncio.create_task(single_client(session, per_client, duration, results)))
        await asyncio.gather(*tasks)
    # compute metrics
    total = len(results)
    succ = sum(1 for ok, _ in results if ok)
    latencies = [lat for ok, lat in results if ok]
    p50 = statistics.median(latencies) if latencies else None
    p95 = sorted(latencies)[int(0.95 * len(latencies))] if latencies else None
    return {"rps": rps, "total": total, "success": succ, "p50": p50, "p95": p95, "latencies": latencies}

def plot_results(points, filename="rps_slo.png", slo_ms=500):
    xs = [p["rps"] for p in points]
    attainment = []
    for p in points:
        lat = p["latencies"]
        if not lat:
            attainment.append(0.0)
        else:
            attainment.append(sum(1 for l in lat if l*1000 <= slo_ms) / len(lat))
    plt.figure()
    plt.plot(xs, attainment, marker="o")
    plt.xlabel("Offered RPS")
    plt.ylabel(f"SLO attainment (TTFT <= {slo_ms} ms)")
    plt.title("RPS vs SLO Attainment")
    plt.grid(True)
    plt.savefig(filename)
    print("Saved plot to", filename)

async def main():
    rps_values = [10, 20, 40, 80, 160]  # adjust based on cluster capacity
    points = []
    for rps in rps_values:
        print("Running RPS", rps)
        p = await run_rps_point(rps=rps, clients=4, duration=20)
        print(p)
        points.append(p)
    plot_results(points)

if __name__ == "__main__":
    asyncio.run(main())
